#ifndef HUFFMANTREEWINDOW_H
#define HUFFMANTREEWINDOW_H

#include <QMainWindow>
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QTextEdit>
#include <QLineEdit>
#include <QPushButton>
#include <QLabel>
#include <QGroupBox>
#include <QMessageBox>
#include <QComboBox>
#include <QTabWidget>
#include <QWidget>
#include <QPainter>
#include <QDialog>
#include <QDialogButtonBox>
#include <QTimer>
#include <QMouseEvent>
#include <QWheelEvent>

struct HuffmanNode {
    int weight;
    char character;
    HuffmanNode* left;
    HuffmanNode* right;
    HuffmanNode(int w, char c = '\0') : weight(w), character(c), left(nullptr), right(nullptr) {}
};

// 比较函数替代std::priority_queue的比较
struct CompareHuffmanNode {
    bool operator()(HuffmanNode* a, HuffmanNode* b) {
        return a->weight > b->weight;
    }
};

// 自定义动态数组替代std::vector
template<typename T>
class SimpleVector {
private:
    T* m_data;
    int m_size;
    int m_capacity;

    void resize(int newCapacity) {
        T* newData = new T[newCapacity];
        for (int i = 0; i < m_size; ++i) {
            newData[i] = m_data[i];
        }
        delete[] m_data;
        m_data = newData;
        m_capacity = newCapacity;
    }

public:
    SimpleVector() : m_data(nullptr), m_size(0), m_capacity(0) {}

    ~SimpleVector() {
        delete[] m_data;
    }

    void push_back(const T& value) {
        if (m_size >= m_capacity) {
            resize(m_capacity == 0 ? 4 : m_capacity * 2);
        }
        m_data[m_size++] = value;
    }

    void pop_back() {
        if (m_size > 0) {
            m_size--;
        }
    }

    int size() const { return m_size; }
    bool empty() const { return m_size == 0; }
    T& operator[](int index) { return m_data[index]; }
    const T& operator[](int index) const { return m_data[index]; }

    void clear() {
        delete[] m_data;
        m_data = nullptr;
        m_size = m_capacity = 0;
    }

    // 添加迭代器支持
    T* begin() { return m_data; }
    T* end() { return m_data + m_size; }
    const T* begin() const { return m_data; }
    const T* end() const { return m_data + m_size; }
};

// 自定义优先队列替代std::priority_queue
template<typename T, typename Compare>
class SimplePriorityQueue {
private:
    SimpleVector<T> m_data;
    Compare m_compare;

    void heapifyUp(int index) {
        while (index > 0) {
            int parent = (index - 1) / 2;
            if (!m_compare(m_data[parent], m_data[index])) break;

            T temp = m_data[parent];
            m_data[parent] = m_data[index];
            m_data[index] = temp;
            index = parent;
        }
    }

    void heapifyDown(int index) {
        int size = m_data.size();
        while (true) {
            int left = 2 * index + 1;
            int right = 2 * index + 2;
            int target = index;

            if (left < size && m_compare(m_data[target], m_data[left])) {
                target = left;
            }
            if (right < size && m_compare(m_data[target], m_data[right])) {
                target = right;
            }
            if (target == index) break;

            T temp = m_data[index];
            m_data[index] = m_data[target];
            m_data[target] = temp;
            index = target;
        }
    }

public:
    SimplePriorityQueue() {}

    void push(const T& value) {
        m_data.push_back(value);
        heapifyUp(m_data.size() - 1);
    }

    void pop() {
        if (empty()) return;
        m_data[0] = m_data[m_data.size() - 1];
        m_data.pop_back();
        if (!empty()) {
            heapifyDown(0);
        }
    }

    T top() const {
        return m_data[0];
    }

    bool empty() const {
        return m_data.empty();
    }

    int size() const {
        return m_data.size();
    }
};

// 自定义映射替代std::map
class CharStringMap {
private:
    struct MapNode {
        char key;
        QString value;
        MapNode* next;
        MapNode(char k, const QString& v) : key(k), value(v), next(nullptr) {}
    };

    static const int TABLE_SIZE = 53;
    MapNode* m_table[TABLE_SIZE];

    int hash(char key) const {
        return static_cast<unsigned char>(key) % TABLE_SIZE;
    }

public:
    CharStringMap() {
        for (int i = 0; i < TABLE_SIZE; ++i) {
            m_table[i] = nullptr;
        }
    }

    ~CharStringMap() {
        clear();
    }

    void insert(char key, const QString& value) {
        int index = hash(key);
        MapNode* current = m_table[index];

        // 检查是否已存在
        while (current) {
            if (current->key == key) {
                current->value = value;
                return;
            }
            current = current->next;
        }

        // 插入新节点
        MapNode* newNode = new MapNode(key, value);
        newNode->next = m_table[index];
        m_table[index] = newNode;
    }

    QString& operator[](char key) {
        int index = hash(key);
        MapNode* current = m_table[index];

        // 查找现有节点
        while (current) {
            if (current->key == key) {
                return current->value;
            }
            current = current->next;
        }

        // 创建新节点
        MapNode* newNode = new MapNode(key, "");
        newNode->next = m_table[index];
        m_table[index] = newNode;
        return newNode->value;
    }

    bool contains(char key) const {
        int index = hash(key);
        MapNode* current = m_table[index];

        while (current) {
            if (current->key == key) {
                return true;
            }
            current = current->next;
        }
        return false;
    }

    void clear() {
        for (int i = 0; i < TABLE_SIZE; ++i) {
            MapNode* current = m_table[i];
            while (current) {
                MapNode* next = current->next;
                delete current;
                current = next;
            }
            m_table[i] = nullptr;
        }
    }

    int count(char key) const {
        return contains(key) ? 1 : 0;
    }
};

// 替代std::pair
struct CharIntPair {
    char first;
    int second;
    CharIntPair() : first('\0'), second(0) {}
    CharIntPair(char f, int s) : first(f), second(s) {}
};

class HuffmanVisualizationWidget : public QWidget
{
    Q_OBJECT

public:
    explicit HuffmanVisualizationWidget(QWidget *parent = nullptr);
    void setRoot(HuffmanNode* rootNode) { root = rootNode; resetView(); update(); }
    void setBuildSteps(const QVector<QVector<HuffmanNode*>>& steps) { buildSteps = steps; totalSteps = steps.size(); currentStep = 0; resetView(); update(); }
    void setCurrentStep(int step) { currentStep = step; update(); }
    int getCurrentStep() const { return currentStep; }
    int getTotalSteps() const { return totalSteps; }

protected:
    void paintEvent(QPaintEvent* event) override;
    void mousePressEvent(QMouseEvent* event) override;
    void mouseMoveEvent(QMouseEvent* event) override;
    void mouseReleaseEvent(QMouseEvent* event) override;
    void wheelEvent(QWheelEvent* event) override;

private:
    void drawTree(QPainter& painter, HuffmanNode* node, int x, int y);
    void drawForest(QPainter& painter, const QVector<HuffmanNode*>& forest, int step);
    int countNodes(HuffmanNode* node);
    void resetView();
    HuffmanNode* root;
    QVector<QVector<HuffmanNode*>> buildSteps;
    int currentStep;
    int totalSteps;

    // 视图变换相关变量
    QPoint lastMousePos;
    bool isDragging;
    QPoint translation;
    double scale;
};

class InputWeightsDialog : public QDialog
{
    Q_OBJECT
public:
    explicit InputWeightsDialog(QWidget *parent = nullptr);
    QString getWeights() const { return weightsEdit->text(); }

private:
    QLineEdit *weightsEdit;
};

class EncodeDecodeDialog : public QDialog
{
    Q_OBJECT
public:
    explicit EncodeDecodeDialog(QWidget *parent = nullptr, bool isEncode = true);
    QString getInput() const { return inputEdit->text(); }

private:
    QLineEdit *inputEdit;
};

class HuffmanTreeWindow : public QMainWindow
{
    Q_OBJECT

public:
    HuffmanTreeWindow(QWidget *parent = nullptr);
    ~HuffmanTreeWindow();

signals:
    void returnToMainMenu();

private slots:
    void onInputWeights();
    void onClear();
    void onGenerateCodes();
    void onEncodeText();
    void onDecodeText();
    void onReturnToMainMenu();
    void onNextStep();
    void onPrevStep();
    void onAutoPlay();

private:
    void setupUI();
    void updateDisplay();
    void clearTree(HuffmanNode* node);
    void buildHuffmanTree();
    void generateHuffmanCodes(HuffmanNode* node, const QString& code);
    int getTreeHeight(HuffmanNode* node);
    void recordBuildStep(const SimpleVector<HuffmanNode*>& nodes);

    // 替换std::max的辅助函数
    int max(int a, int b) { return a > b ? a : b; }

    HuffmanNode* m_root;
    CharStringMap m_huffmanCodes;
    SimpleVector<CharIntPair> m_frequencies;

    QVector<QVector<HuffmanNode*>> m_buildSteps;
    int m_currentBuildStep;

    QTextEdit *m_display;
    HuffmanVisualizationWidget *m_visualizationWidget;
    QPushButton *m_nextStepBtn;
    QPushButton *m_prevStepBtn;
    QPushButton *m_autoPlayBtn;
    QTimer *m_autoPlayTimer;
};

#endif // HUFFMANTREEWINDOW_H
